Hugging FaceでEarlyStoppingを使用する方法

nokomoro3

2022.11.04

この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。

こんちには。

データアナリティクス事業本部機械学習チームの中村です。

Hugging Faceのライブラリの使い方紹介記事第２弾です。

今回は、学習時にEarly Stoppingを使用する方法をご紹介します。

Early Stoppingとは

Deep Learningなどのモデルを学習する際には、通常trainデータとvalidデータ(検定データ)に分割し、 trainデータでモデル自体を更新し、validデータで学習外データに対する汎化性能を確認します。

より具体的にはvalidデータの損失関数などの指標を確認し、最良となるようなepochで学習を停止するなどの操作をします。これがEarly Stoppingと一般的に呼ばれるものです。

実行環境

今回はGoogle Colaboratory環境で実行しました。

ハードウェアなどの情報は以下の通りです。

GPU: Tesla P100 (GPUメモリ16GB搭載)
CUDA: 11.1
メモリ: 13GB

主なライブラリのバージョンは以下となります。

transformers: 4.22.1
datasets: 2.4.0

インストール

transformersとdatasetsをインストールします。

!pip install transformers datasets

また事前学習モデルの依存モジュールをインストールします。

!pip install fugashi
!pip install ipadic
!pip install sentencepiece

ベースとするコード

今回のベースとするコードは以下のとおりです。

from datasets import load_dataset
from transformers import AutoTokenizer
from transformers import AutoModelForSequenceClassification
from transformers import TrainingArguments
from transformers import Trainer
from sklearn.metrics import accuracy_score, f1_score
import torch

# データセットのロード
dataset = load_dataset("tyqiangz/multilingual-sentiments", "japanese")

# # 実験のためデータセットを縮小したい場合はコチラを有効化
# dataset = DatasetDict({
#     "train": dataset['train'].select(range(1000)),
#     "validation": dataset['validation'].select(range(1000)),
#     "test": dataset['test'].select(range(1000)),
# })

# トークナイザのロード
model_ckpt = "cl-tohoku/bert-base-japanese-whole-word-masking"
tokenizer = AutoTokenizer.from_pretrained(model_ckpt)

# トークナイズ処理
def tokenize(batch):
    return tokenizer(batch["text"], padding=True, truncation=True)
dataset_encoded = dataset.map(tokenize, batched=True, batch_size=None)

# 事前学習モデルのロード
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
num_labels = 3
model = (AutoModelForSequenceClassification
    .from_pretrained(model_ckpt, num_labels=num_labels)
    .to(device))

# メトリクスの定義
def compute_metrics(pred):
    labels = pred.label_ids
    preds = pred.predictions.argmax(-1)
    f1 = f1_score(labels, preds, average="weighted")
    acc = accuracy_score(labels, preds)
    return {"accuracy": acc, "f1": f1}

# 学習パラメータの設定
batch_size = 16
logging_steps = len(dataset_encoded["train"]) // batch_size
model_name = "sample-text-classification-bert"

training_args = TrainingArguments(
    output_dir=model_name,
    num_train_epochs=2,
    learning_rate=2e-5,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    weight_decay=0.01,
    evaluation_strategy="epoch",
    disable_tqdm=False,
    logging_steps=logging_steps,
    push_to_hub=False,
    log_level="error",
)

# Trainerの定義
trainer = Trainer(
    model=model,
    args=training_args,
    compute_metrics=compute_metrics,
    train_dataset=dataset_encoded["train"],
    eval_dataset=dataset_encoded["validation"],
    tokenizer=tokenizer
)

# トレーニング実行
trainer.train()

この内容についての解説は以下の記事を参照ください。

TrainingArgumentの変更

EarlyStoppingを使用するには、基本的にはコールバックを設定することになりますですが、使用するためにはTrainingArgumentが特定の条件を満たすことが必要です。

今回は、以下の２点をTrainingArgumentで変更しました。

load_best_model_at_end=Trueとする必要があったため指定
evaluation_strategyとsave_strategyが同じである必要があったため指定
- 正確には、これはload_best_model_at_end=Trueとする際の条件

またnum_train_epochs=2ではEarlyStoppingが働く前に学習が終わりますので、num_train_epochs=20に変更しています。

training_args = TrainingArguments(
    output_dir=model_name,
    num_train_epochs=20,
    learning_rate=2e-5,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    weight_decay=0.01,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    disable_tqdm=False,
    logging_steps=logging_steps,
    push_to_hub=False,
    log_level="error",
    load_best_model_at_end=True,
)

EarlyStoppingコールバックを指定

EarlyStoppingのコールバックをTrainerクラスのcallbacksに指定します。

early_stopping_patienceは、どの程度eval_datasetで計算される指標が改善しなかったら学習を停止するかの設定で、今回は3(この場合単位はエポックとなる)としています。

from transformers import EarlyStoppingCallback

trainer = Trainer(
    model=model,
    args=training_args,
    compute_metrics=compute_metrics,
    train_dataset=dataset_encoded["train"],
    eval_dataset=dataset_encoded["validation"],
    tokenizer=tokenizer,
    callbacks=[EarlyStoppingCallback(early_stopping_patience=3)],
)

trainer.train()

なお、ベストと判断するための指標は、デフォルトではlossが使用されます。

別の指標（例えばF1スコアなど）を判断基準にしたい場合は、metric_for_best_model="f1"などと指定すればOKです。この場合、本記事の例のようにcompute_metricsを同じ名前("f1")をreturnするように定義して与える必要があります。